package com.example.guance_demo_emqx;

import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;

import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import org.eclipse.paho.mqttv5.client.MqttClient;
import org.eclipse.paho.mqttv5.client.MqttConnectionOptions;
import org.eclipse.paho.mqttv5.client.MqttCallback;
import org.eclipse.paho.mqttv5.client.IMqttToken;
import org.eclipse.paho.mqttv5.client.MqttDisconnectResponse;
import org.eclipse.paho.mqttv5.common.MqttMessage;
import org.eclipse.paho.mqttv5.common.MqttException;
import org.eclipse.paho.mqttv5.common.packet.MqttProperties;
import org.eclipse.paho.mqttv5.common.packet.UserProperty;

// Datadog 扩展了 Opentracing 的标签，两个包中的字段共同组成了 Datadog 的追踪数据协议
import datadog.trace.api.DDTags;
import io.opentracing.tag.Tags;
import io.opentracing.Tracer;
import io.opentracing.noop.NoopSpan;
import io.opentracing.Span;
import io.opentracing.SpanContext;
import io.opentracing.Scope;
import io.opentracing.util.GlobalTracer;
import io.opentracing.propagation.TextMapAdapter;
import io.opentracing.propagation.Format.Builtin;
 
@Component
public class MqttClientService {
 
    @Value("${mqtt.server.host}")
    private String mqttServerHost;
 
    @Value("${mqtt.server.username}")
    private String mqttUsername;
 
    @Value("${mqtt.server.password}")
    private String mqttPassword;
 
    private MqttClient mqttClient;
 
    @PostConstruct
    public void init() {
        try {
            String clientId = String.join("_", "mqttx", UUID.randomUUID().toString().replace("-", ""));
            mqttClient = new MqttClient(mqttServerHost, clientId);
 
            MqttConnectionOptions options = new MqttConnectionOptions();
            options.setUserName(mqttUsername);
            options.setPassword(mqttPassword.getBytes());
 
            mqttClient.connect(options);

            mqttClient.subscribe("topic/test", 1);
            mqttClient.setCallback(new MqttCallback() {
                @Override
                public void messageArrived(String topic, MqttMessage message) throws Exception {
                    // 对于 Opentracing，正常情况下应该通过适配某种协议的适配器获取追踪信息来构建一个 Span 上下文，新的 Span 继
                    // 承该上下文实现链路串联，此场景下，不实现专用的适配器，而是使用通用适配器，手动增加 traceparent 字段，由于
                    // Datadog 探针兼容 W3C 链路传播协议，能够实现生产者和消费者的消息串联

                    // 注意：traceparent 作为用户属性，在生产者发送和消费者接收期间保持不变，这是符合 MQTTv5 规范的，但是这也就
                    // 造成了 EMQX 内部拆分的 Span 无法与消费者关联的情况，这个问题在项目 Issue 中有讨论：https://github.com/emqx/emqx/issues/12752

                    MqttProperties properties = message.getProperties();

                    // 从用户属性中获取 traceparent
                    String traceparent = "";
                    if (properties != null) {
                        List<UserProperty> userProperties = properties.getUserProperties();
                        if (userProperties != null) {
                            for (UserProperty userProperty : userProperties) {
                                if ("traceparent".equals(userProperty.getKey())) {
                                    traceparent = userProperty.getValue();
                                }
                            }
                        }
                    }

                    // 如果获取到符合规范的 traceparent 则创建上下文，否则直接消费
                    if (traceparent != "" && FormatValidator.validateTraceparentFormat(traceparent)) {
                        Tracer tracer = GlobalTracer.get();
                        
                        // 使用通用适配器创建 Span 上下文
                        Map<String, String> context = new HashMap<>();
                        context.put("traceparent", traceparent);
                        SpanContext spanContext = tracer.extract(Builtin.TEXT_MAP, new TextMapAdapter(context));

                        // 新的 Span 继承上下文
                        final Span span = tracer.buildSpan("mqtt.subscribe")
                            .asChildOf(spanContext)
                            .withTag(DDTags.SERVICE_NAME, "guance-demo-emqx")
                            .withTag(DDTags.RESOURCE_NAME, "GuanceDemoEmqxApplication.subscribeMessage")
                            .start();

                        try (Scope scope = tracer.activateSpan(span)) {
                            // 丰富标签
                            span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_CONSUMER);

                            // 模拟消息消费
                            Thread.sleep(10);
                            System.out.println("[Trace] Message arrived. Topic: " + topic + " Payload: " + new String(message.getPayload()));
                        } catch (Exception e) {
                            // 为 Span 设置错误
                            span.setTag(Tags.ERROR, true);
                            span.setTag(DDTags.ERROR_MSG, e.getMessage());
                            span.setTag(DDTags.ERROR_TYPE, e.getClass().getName());

                            final StringWriter errorString = new StringWriter();
                            e.printStackTrace(new PrintWriter(errorString));
                            span.setTag(DDTags.ERROR_STACK, errorString.toString());
                        } finally {
                            span.finish();
                        }
                    } else {
                        System.out.println("[No Trace] Message arrived. Topic: " + topic + " Payload: " + new String(message.getPayload()));
                    }
                }

                @Override
                public void deliveryComplete(IMqttToken token){}

                @Override
                public void authPacketArrived(int reasonCode, MqttProperties properties) {}

                @Override
                public void connectComplete(boolean reconnect, String serverUR) {}

                @Override
                public void disconnected(MqttDisconnectResponse disconnectResponse) {}

                @Override
                public void mqttErrorOccurred(MqttException exception) {}
            });
        } catch (MqttException e) {
            e.printStackTrace();
        }
    }
 
    @PreDestroy
    public void cleanup() {
        try {
            mqttClient.disconnect();
            mqttClient.close();
        } catch (MqttException e) {
            e.printStackTrace();
        }
    }
 
    public void publishMessage(String topic, String payload) {
        Tracer tracer = GlobalTracer.get();
        // 获取当前活动的 Span
        final Span span = tracer.activeSpan();
        String traceparent = "";
        if (span != null && !(span instanceof NoopSpan)) {
            System.out.println(span);
            String traceId = span.context().toTraceId();
            String spanId = span.context().toSpanId();

            System.out.println("Trace ID: " + traceId);
            System.out.println("Original Span ID: " + spanId);

            // 构造 traceparent
            // 注意：
            //   1. Datadog SpanId 默认使用十进制，与 OTel 集成时需转换为十六进制，否则 EMQX 无法创建追踪；
            //   2. Datadog SpanId 转换为十六进制后字符串长度可能小于 16，需要前补 0，否则 EMQX 无法创建追踪；
            //   3. 观测云 Datakit 采集器支持统一转换 SpanId 为十进制或者十六进制，以实现链路关联。
            spanId = Long.toHexString(Long.parseLong(spanId));
            spanId = String.format("%16s", spanId).replace(' ', '0');

            // W3C 标准：version - trace-id - parent-id - trace-flags
            traceparent = "00-" + traceId + "-" + spanId + "-01";

            System.out.println("Processed Span ID: " + new String(spanId));
        }

        MqttProperties mqttProperties = new MqttProperties();
        List<UserProperty> userProperties = new ArrayList<>();

        if (!traceparent.isEmpty()) {
            userProperties.add(new UserProperty("traceparent", traceparent));
        }

        mqttProperties.setUserProperties(userProperties);

        try {
            // 模拟一个延迟，以命中 Profile 采样
            Thread.sleep(1000);

            MqttMessage message = new MqttMessage(payload.getBytes());
            // 为消息设置用户属性
            message.setProperties(mqttProperties);
            mqttClient.publish(topic, message);
        } catch (InterruptedException e) {
            e.printStackTrace();
        } catch (MqttException e) {
            e.printStackTrace();
        }
    }
}
